#include <stdio.h>
#include <string.h>
#include <assert.h>
#include "pkcs11.h"
#include "session.h"

void	test_enc(CK_SESSION_HANDLE hSession)
{
	CK_RV ret;
	CK_MECHANISM m;
	CK_OBJECT_HANDLE key;
	CK_BYTE part[256] = {"msg part"};
	CK_ULONG partlen = 8;
	CK_BYTE	out[256];
	CK_ULONG outlen;

	m.mechanism = 10;
	m.ulParameterLen = 3;
	m.pParameter = "zcp";
	key = 12345678;
	ret = C_EncryptInit(hSession, &m, key);
	assert(ret == CKR_OK);

	memset(out, 0, 256);
	ret = C_EncryptUpdate(hSession, part, partlen, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);

	memset(out, 0, 256);
	ret = C_EncryptFinal(hSession, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);

	memset(out, 0, 256);
	ret = C_Encrypt(hSession, part, partlen, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);
}

void	test_dec(CK_SESSION_HANDLE hSession)
{
	CK_RV ret;
	CK_MECHANISM m;
	CK_OBJECT_HANDLE key;
	CK_BYTE part[256] = {"msg de part"};
	CK_ULONG partlen = 11;
	CK_BYTE	out[256];
	CK_ULONG outlen;

	m.mechanism = 11;
	m.ulParameterLen = 4;
	m.pParameter = "zcpa";
	key = 12345679;
	ret = C_DecryptInit(hSession, &m, key);
	assert(ret == CKR_OK);

	memset(out, 0, 256);
	ret = C_DecryptUpdate(hSession, part, partlen, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);

	memset(out, 0, 256);
	ret = C_DecryptFinal(hSession, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);

	memset(out, 0, 256);
	ret = C_Decrypt(hSession, part, partlen, out, &outlen);
	assert(ret == CKR_OK);
	printf("out: %s\n", out);
	printf("outlen: %ld\n", outlen);
}

void	test_digest(CK_SESSION_HANDLE hSession)
{
	CK_RV ret;
	CK_MECHANISM m;
	CK_BYTE data[256] = {"abcccc"}, out[256];
	CK_ULONG len =  6, outlen;

	m.mechanism = 12;
	m.ulParameterLen = 5;
	m.pParameter = "zcp12";
	ret = C_DigestInit(hSession, &m);
	assert(ret == CKR_OK);

	ret = C_DigestUpdate(hSession, data, len);
	assert(ret == CKR_OK);

	memset(out, 0, 256);
	ret = C_DigestFinal(hSession, out, &outlen);
	assert(ret == CKR_OK);
	printf("outlen : %ld\n", outlen);
	printf("out: %s\n", (char *)out);

	memset(out, 0, 256);
	ret = C_Digest(hSession, data, len, out, &outlen); 
	assert(ret == CKR_OK);
	printf("outlen : %ld\n", outlen);
	printf("out: %s\n", (char *)out);
}

void	test_sign_verify(CK_SESSION_HANDLE hSession)
{
	CK_RV	ret;
	CK_BYTE	msg[256] = {"signmsg"}, sig[256] = {0};
	CK_ULONG msglen = 7, siglen;

	ret = C_Sign(hSession, msg, msglen, sig, &siglen);
	assert(ret == CKR_OK);
	printf("sign len : %ld\n", siglen);
	printf("sign: %s\n", (char *)sig);

	ret = C_Verify(hSession, msg, msglen, sig, siglen);
	assert(ret == CKR_OK);
}

void	test_key_manage(CK_SESSION_HANDLE hSession)
{
	CK_RV	ret, i;
	CK_MECHANISM m;
	CK_ATTRIBUTE temp[3], pub_temp[3], pri_temp[3];
	CK_ULONG count = 3, len;
	CK_OBJECT_HANDLE key, key2, pubkey, prikey;
	CK_BYTE data[256];
	char str[3][256];
	char str2[3][256];

	m.mechanism = 13;
	m.ulParameterLen = 4;
	m.pParameter = "keym";

	for(i = 0; i < count; i++)
	{
		sprintf(str[i], "key_temp%ld", i);
		temp[i].type = 88+i;
		temp[i].pValue = str[i]; 
		temp[i].ulValueLen = strlen(str[i]);
	}
	ret = C_GenerateKey(hSession, &m, temp, count, &key);
	assert(ret == CKR_OK);
	printf("key: %ld\n", key);

	m.mechanism = 14;
	m.ulParameterLen = 7;
	m.pParameter = "keypair";

	for(i = 0; i < count; i++)
	{
		sprintf(str[i], "pubkey_temp%ld", i);
		pub_temp[i].type = 8+i;
		pub_temp[i].pValue = str[i]; 
		pub_temp[i].ulValueLen = strlen(str[i]);
	}
	for(i = 0; i < count; i++)
	{
		sprintf(str2[i], "prikey_temp%ld", i);
		pri_temp[i].type = 18 + i;
		pri_temp[i].pValue = str2[i]; 
		pri_temp[i].ulValueLen = strlen(str2[i]);
	}

	ret = C_GenerateKeyPair(hSession, &m, pub_temp, count, pri_temp, count, &pubkey, &prikey);
	assert(ret == CKR_OK);
	printf("pub key: %ld\n", pubkey);
	printf("pri key: %ld\n", prikey);

	for(i = 0; i < 3; i++)
	{
		sprintf(str[i], "keyc_temp%ld", i);
		temp[i].type = 28+i;
		temp[i].pValue = str[i]; 
		temp[i].ulValueLen = strlen(str[i]);
	}
	ret = C_CreateObject(hSession, temp, 3, &key); 
	assert(ret == CKR_OK);
	printf("key: %ld\n", key);

	for(i = 0; i < 3; i++)
	{
		sprintf(str[i], "setattr_temp%ld", i);
		temp[i].type = 38+i;
		temp[i].pValue = str[i]; 
		temp[i].ulValueLen = strlen(str[i]);
	}
	ret = C_SetAttributeValue(hSession, key, temp, 3);
	assert(ret == CKR_OK);

	memset(data, 0, 256);
	m.mechanism = 15;
	ret = C_WrapKey(hSession, &m, prikey, key, data, &len); 
	assert(ret == CKR_OK);
	printf("wraped: %s\n", (char *)data);
	printf("wraped len: %ld\n", len);

	m.mechanism = 16;
	ret = C_UnwrapKey(hSession, &m, key, data, len, temp, 3, &key2);
	assert(ret == CKR_OK);
	printf("C_UnwrapKey: %ld\n", key2);

#if 1
	char temp_str0[256];
	char temp_str1[256];
	char temp_str2[256];

	memset(&temp[0], 0, sizeof(CK_ATTRIBUTE));
	memset(&temp[1], 0, sizeof(CK_ATTRIBUTE));
	memset(&temp[2], 0, sizeof(CK_ATTRIBUTE));
	temp[0].type = 6600;
	temp[0].pValue = temp_str0;
	temp[0].ulValueLen = 256; 

	temp[1].type = 6601;
	temp[1].pValue = temp_str1;
	temp[1].ulValueLen = 256; 

	temp[2].type = 6602;
	temp[2].pValue = temp_str2;
	temp[2].ulValueLen = 256; 

	ret = C_GetAttributeValue(hSession, key, temp, 3);
	assert(ret == CKR_OK);
	for(i = 0; i < 3; i++)
	{
		char attr_str[256];

		memset(attr_str, 0, 256);
		memcpy(attr_str, temp[i].pValue, temp[i].ulValueLen);
		printf("get attrvalue: %s\n", attr_str);
	}
#endif

	for(i = 0; i < 3; i++)
	{
		sprintf(temp[i].pValue, "setattr%ld", i);
		temp[i].ulValueLen = strlen(temp[i].pValue);
	}
	ret = C_SetAttributeValue(hSession, key, temp, 3);
	assert(ret == CKR_OK);

	ret = C_DestroyObject(hSession, key); 
	assert(ret == CKR_OK);
}

void	test_find(CK_SESSION_HANDLE hSession)
{
	CK_RV ret, i;
	CK_ATTRIBUTE temp[3];
	CK_OBJECT_HANDLE keys[128];
	CK_ULONG num;
	char str[3][256];

	for(i = 0; i < 3; i++)
	{
		sprintf(str[i], "find_temp%ld", i);
		temp[i].type = 68+i;
		temp[i].pValue = str[i];
		temp[i].ulValueLen = strlen(str[i]);
	}
	ret = C_FindObjectsInit(hSession, temp, 3);
	assert(ret == CKR_OK);

	ret = C_FindObjects(hSession, 128, keys, &num);
	assert(ret == CKR_OK);

	printf("find num: %ld\n", num);
	for(i = 0; i < num; i++)
	{
		printf("find obj: %ld\n", keys[i]);
	}

	ret = C_FindObjectsFinal(hSession);
	assert(ret == CKR_OK);
}

int	main()
{
	CK_RV ret;
	CK_BYTE r[128] = {0};
	CK_SESSION_HANDLE hSession;

	ret = C_OpenSession(0, 0, 0, 0, &hSession); 
	assert(ret == CKR_OK);

	ret = C_LoginUser(hSession, 1, (CK_UTF8CHAR_PTR)"123456", 6, (CK_UTF8CHAR_PTR)"forxy", 5);
	assert(ret == CKR_OK);

	ret = C_GenerateRandom(hSession, r, 32);
	assert(ret == CKR_OK);
	printf("random : %s\n", r); 

	test_enc(hSession);
	test_dec(hSession);
	test_digest(hSession);
	test_sign_verify(hSession);
	test_key_manage(hSession);
	test_find(hSession);

	ret = C_Logout(hSession);
	assert(ret == CKR_OK);

	printf("test ok.\n");
	return 0;
}
